import tensorflow as tf
import os
from tensorflow.python.tools import freeze_graph
from C3D_model import C3D as network
 
model_path = "check_point/train.ckpt-24" 

BATCH_SIZE = 1
#BATCH_SIZE = 1
#NUM_CLASSES = 101
NUM_CLASSES = 10
CROP_SZIE = 112
CHANNEL_NUM = 3
CLIP_LENGTH = 16

def main():
    tf.reset_default_graph()
    
    batch_clips = tf.placeholder(tf.float32, [BATCH_SIZE, CLIP_LENGTH, CROP_SZIE, CROP_SZIE, CHANNEL_NUM], name='X')
    #batch_labels = tf.placeholder(tf.int32, [BATCH_SIZE, NUM_CLASSES], name='Y')
    #keep_prob = tf.placeholder(tf.float32)

    #input_node = tf.placeholder(tf.float32, shape=(None, None, 200)) 
    #num_classes = 2
    output_node = network(batch_clips) 
    
    flow = tf.cast(output_node , tf.float16, 'the_outputs') 
    saver = tf.train.Saver()
    with tf.Session() as sess:
        saver.restore(sess, model_path)
    
    tf.train.write_graph(sess.graph_def, 'output_model/pb_model', 'model.pb')
    
    freeze_graph.freeze_graph('output_model/pb_model/model.pb', '', False, model_path, 'the_outputs','save/restore_all', 'save/Const:0', 'output_model/pb_model/frozen_model.pb', False, "")
    print("done")

if __name__ == '__main__':
    main()
